//#include <stdexcept>
#include <atomic>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <dlfcn.h>
#include <cxxabi.h>
#include <fcntl.h>
#include <string.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <pthread.h>
#include <signal.h>

#define UNW_LOCAL_ONLY
#include <libunwind.h>

#define MAX_BT_DEPTH 32
#define MAX_SYMBOL_LEN 64

extern "C" void *__libc_malloc(size_t size);
extern "C" void __libc_free(void *addr);

void signal_handler(int signum);

// 哈希表节点结构体
struct HashNode
{
    void*       key;
    char**      value;
    HashNode*   next;
};

// 哈希表类
class HashTable
{
private:
    // 哈希函数
    unsigned int hash(void *key)
    {
        return reinterpret_cast<unsigned long>(key) % m_tableSize;
    }

    // rehash 函数，用于扩容或缩容后重新分配元素
    void rehash(int newSize)
    {
        HashNode **oldTable = m_table;
        int oldSize = m_tableSize;
        m_tableSize = newSize;
        m_table = static_cast<HashNode **>(__libc_malloc(m_tableSize * sizeof(HashNode *)));
        for (int i = 0; i < m_tableSize; ++i)
        {
            m_table[i] = nullptr;
        }
        m_elementCount = 0;

        for (int i = 0; i < oldSize; ++i)
        {
            HashNode *current = oldTable[i];
            while (current != nullptr)
            {
                insert(current->key, current->value);
                current = current->next;
            }
        }
        __libc_free(oldTable);
    }

public:
    HashTable(int initialSize = 10) : m_tableSize(initialSize), m_elementCount(0)
    {
        printf("HashTable::HashTable()\n");

        pthread_mutexattr_t attr;
        pthread_mutexattr_init(&attr);
        pthread_mutexattr_settype(&attr, PTHREAD_MUTEX_RECURSIVE);//设置递归锁，避免insert->rehash->insert死锁
        pthread_mutex_init(&m_mutex, &attr);

        m_table = static_cast<HashNode **>(__libc_malloc(m_tableSize * sizeof(HashNode *)));
        for (int i = 0; i < m_tableSize; ++i)
        {
            m_table[i] = nullptr;
        }
    }

    ~HashTable()
    {
        for (int i = 0; i < m_tableSize; ++i)
        {
            HashNode *current = m_table[i];
            while (current != nullptr)
            {
                HashNode *temp = current;
                current = current->next;

                // 释放字符串数组内存
                for (int j = 0; temp->value[j] != nullptr; ++j)
                {
                    __libc_free(temp->value[j]);
                }
                __libc_free(temp->value);

                __libc_free(temp);
            }
        }
        __libc_free(m_table);

        pthread_mutex_destroy(&m_mutex);
    }

    // 插入键值对，value最后一个char*为nullptr标志结束
    void insert(void *key, char **value)
    {
        pthread_mutex_lock(&m_mutex);
        if (static_cast<double>(m_elementCount) / m_tableSize >= LOAD_FACTOR_THRESHOLD_UP)
        {
            rehash(m_tableSize * 2);
        }

        unsigned int index = hash(key);
        HashNode *newNode = static_cast<HashNode *>(__libc_malloc(sizeof(HashNode)));
        newNode->key = key;
        newNode->value = value;
        newNode->next = nullptr;

        if (m_table[index] == nullptr)
        {
            m_table[index] = newNode;
        }
        else
        {
            HashNode *current = m_table[index];
            while (current->next != nullptr)
            {
                current = current->next;
            }
            current->next = newNode;
        }
        ++m_elementCount;
        pthread_mutex_unlock(&m_mutex);
    }

    // 根据键查找值
    char **get(void *key)
    {
        pthread_mutex_lock(&m_mutex);
        unsigned int index = hash(key);
        HashNode *current = m_table[index];

        while (current != nullptr)
        {
            if (current->key == key)
            {
                pthread_mutex_unlock(&m_mutex);
                return current->value;
            }
            current = current->next;
        }
        pthread_mutex_unlock(&m_mutex);
        return nullptr;
    }

    // 删除键值对
    void remove(void *key)
    {
        pthread_mutex_lock(&m_mutex);
        unsigned int index = hash(key);
        HashNode *current = m_table[index];
        HashNode *prev = nullptr;

        while (current != nullptr && current->key != key)
        {
            prev = current;
            current = current->next;
        }

        if (current == nullptr)
        {
            pthread_mutex_unlock(&m_mutex);
            return;
        }

        if (prev == nullptr)
        {
            m_table[index] = current->next;
        }
        else
        {
            prev->next = current->next;
        }

        // 释放字符串数组内存
        for (int i = 0; current->value[i] != nullptr; ++i)
        {
            __libc_free(current->value[i]);
        }
        __libc_free(current->value);

        __libc_free(current);
        --m_elementCount;

        if (m_tableSize > 10 && static_cast<double>(m_elementCount) / m_tableSize <= LOAD_FACTOR_THRESHOLD_DOWN)
        {
            rehash(m_tableSize / 2);
        }
        pthread_mutex_unlock(&m_mutex);
    }

    // 获取当前哈希表的大小
    int getTableSize() const
    {
        return m_tableSize;
    }

    //获取所有元素，最后一个HashNode*为nullptr标志结束
    HashNode** getElements()
    {
        HashNode** elements = static_cast<HashNode**>(__libc_malloc((m_elementCount + 1) * sizeof(HashNode*)));

        pthread_mutex_lock(&m_mutex);
        for(int i = 0, j = 0; i < m_tableSize; ++i)
        {
            HashNode* current = m_table[i];
            while(current)
            {
                HashNode* newNode = static_cast<HashNode*>(__libc_malloc(sizeof(HashNode)));

                //拷贝字符串数组
                char** newValue = static_cast<char**>(__libc_malloc((MAX_BT_DEPTH + 1) * sizeof(char*)));
                memset(newValue, 0, (MAX_BT_DEPTH + 1) * sizeof(char*));
                for(int k = 0; current->value[k] != nullptr && k < MAX_BT_DEPTH; ++k)
                {
                    newValue[k] = static_cast<char*>(__libc_malloc(MAX_SYMBOL_LEN));
                    memcpy(newValue[k], current->value[k], MAX_SYMBOL_LEN);
                }

                newNode->key = current->key;
                newNode->value = newValue;
                newNode->next = nullptr;
                elements[j++] = newNode;
                current = current->next;
            }
        }
        pthread_mutex_unlock(&m_mutex);

        elements[m_elementCount] = nullptr;
        return elements;
    }

private:
    pthread_mutex_t m_mutex;
    HashNode**      m_table;
    std::atomic_int m_tableSize;
    std::atomic_int m_elementCount;
    const double LOAD_FACTOR_THRESHOLD_UP = 0.7;
    const double LOAD_FACTOR_THRESHOLD_DOWN = 0.2;
};

bool g_isProcInit = false;

__attribute__((constructor)) void preInitialization() 
{
    printf("preInitialization()\n");
    g_isProcInit = true;
}

// 内存检查助手类
class MemChkHelper
{
public:
    static MemChkHelper& getInstance()
    {
        printf("MemChkHelper::getInstance()\n");
        static MemChkHelper instance;
        return instance;
    }

    ~MemChkHelper()
    {
        record();
        printf("MemChkHelper::~MemChkHelper()\n");
    }

    void record()
    {
        int fd = open("memchk.log", O_WRONLY | O_CREAT | O_TRUNC, 0644);
        if(fd == -1)
        {
            //throw std::runtime_error("open memchk.log failed!");
            printf("open memchk.log failed!\n");
            return;
        }

        char buf[256];

        HashNode** elements = m_mapMalloc.getElements();
        for(int i = 0; elements[i] != nullptr; ++i)
        {
            memset(buf, 0, 256);

            void* addr = elements[i]->key;
            char** attr = m_mapAttributes.get(addr);
            if(!attr)
            {
                printf("addr=%p, attr is nullptr!\n", addr);
                return;
            }

            char* strLen = attr[0];
            snprintf(buf, 256, "malloc(size=%s), addr=%p\n", strLen, addr);
            write(fd, buf, strlen(buf));
            
            char** bt = elements[i]->value;
            for(int idx = 0; bt[idx] != nullptr; ++idx)
            {
                memset(buf, 0, 256);

                char* sym = bt[idx];
                if(strcmp(sym, "??") == 0)
                {
                    snprintf(buf, 256, "#%d ??\n", idx);
                    write(fd, buf, strlen(buf));
                }
                else
                {
                    char* cxxsym = abi::__cxa_demangle(sym, NULL, NULL, NULL);
                    snprintf(buf, 256, "#%d %s >> %s\n", idx, sym, cxxsym);
                    write(fd, buf, strlen(buf));
                }
            }
            write(fd, "\n", 1);
        }

        elements = m_mapNew.getElements();
        for(int i = 0; elements[i]!= nullptr; ++i)
        {
            memset(buf, 0, 256);

            void* addr = elements[i]->key;
            char** attr = m_mapAttributes.get(addr);
            if(!attr)
            {
                printf("addr=%p, attr is nullptr!\n", addr);
                return;
            }

            char* strLen = attr[0];
            snprintf(buf, 256, "new(size=%s), addr=%p\n", strLen, addr);
            write(fd, buf, strlen(buf));

            char** bt = elements[i]->value;
            for(int idx = 0; bt[idx]!= nullptr; ++idx)
            {
                memset(buf, 0, 256);
                char* sym = bt[idx];
                if(strcmp(sym, "??") == 0)
                {
                    snprintf(buf, 256, "#%d ??\n", idx);
                    write(fd, buf, strlen(buf));
                }
                else
                {
                    char* cxxsym = abi::__cxa_demangle(sym, NULL, NULL, NULL);
                    snprintf(buf, 256, "#%d %s >> %s\n", idx, sym, cxxsym);
                    write(fd, buf, strlen(buf));
                }
            }
            write(fd, "\n", 1);
        }

        close(fd);
    }

    //添加malloc记录，bt为堆栈符号列表, attr为属性列表
    void addMallocInfo(void *addr, char **bt, char** attr = nullptr)
    {
        if (!addr)
        {
            //throw std::runtime_error("add_mem: invalid address!");
            return;
        }
        if (!bt)
        {
            //throw std::runtime_error("add_mem: invalid backtrace!");
            return;
        }
        if (m_mapMalloc.get(addr))
        {
            //throw std::runtime_error("add_mem: duplicate address!");
            return;
        }

        m_mapMalloc.insert(addr, bt);

        if(attr)
        {
            m_mapAttributes.insert(addr, attr);
        }
    }

    //删除malloc记录
    void removeMallocInfo(void *addr)
    {
        if (!addr)
        {
            //throw std::runtime_error("remove_mem: invalid address!");
            return;
        }
        if (!m_mapMalloc.get(addr))
        {
            //throw std::runtime_error("remove_mem: unknow address!");
            return;
        }

        m_mapMalloc.remove(addr);
        m_mapAttributes.remove(addr);
    }

    //添加new记录，bt为堆栈符号列表, attr为属性列表
    void addNewInfo(void *addr, char **bt, char** attr = nullptr)
    {
        if (!addr)
        {
            //throw std::runtime_error("add_mem: invalid address!");
            return;
        }
        if (!bt)
        {
            //throw std::runtime_error("add_mem: invalid backtrace!");
            return;
        }
        if (m_mapNew.get(addr))
        {
            //throw std::runtime_error("add_mem: duplicate address!");
            return;
        }

        m_mapNew.insert(addr, bt);
        if(attr)
        {
            m_mapAttributes.insert(addr, attr);
        }
    }

    //删除new记录
    void removeNewInfo(void *addr)
    {
        if (!addr)
        {
            //throw std::runtime_error("remove_mem: invalid address!");
            return;
        }
        if (!m_mapNew.get(addr))
        {
            //throw std::runtime_error("remove_mem: unknow address!");
            return;
        }

        m_mapNew.remove(addr);
        m_mapAttributes.remove(addr);
    }

private:
    MemChkHelper()
    {
        printf("MemChkHelper::MemChkHelper()\n");

        //注册信号处理函数`
        if (signal(SIGTERM, signal_handler) == SIG_ERR) 
        {
            perror("Failed to register SIGTERM handler");
        }
        if (signal(SIGINT, signal_handler) == SIG_ERR) 
        {
            perror("Failed to register SIGINT handler");
        }
        if (signal(SIGHUP, signal_handler) == SIG_ERR) 
        {
            perror("Failed to register SIGHUP handler");
        }
        if (signal(SIGQUIT, signal_handler) == SIG_ERR) 
        {
            perror("Failed to register SIGQUIT handler");
        }
        if (signal(SIGABRT, signal_handler) == SIG_ERR)
        {
            perror("Failed to register SIGABRT handler");
        }
        if (signal(SIGILL, signal_handler) == SIG_ERR)
        {
            perror("Failed to register SIGILL handler");
        }
        if (signal(SIGFPE, signal_handler) == SIG_ERR)
        {
            perror("Failed to register SIGFPE handler");
        }
        if (signal(SIGSEGV, signal_handler) == SIG_ERR)
        {
            perror("Failed to register SIGSEGV handler");
        }
    }

private:
    HashTable       m_mapMalloc;
    HashTable       m_mapNew;
    HashTable       m_mapAttributes;
};

// 信号处理函数
void signal_handler(int signum) 
{
    switch (signum) {
        case SIGTERM:
            printf("Received SIGTERM, cleaning up and exiting...\n");
            MemChkHelper::getInstance().record();
            break;
        case SIGINT:
            printf("Received SIGINT, cleaning up and exiting...\n");
            MemChkHelper::getInstance().record();
            break;
        case SIGHUP:
            printf("Received SIGHUP, cleaning up and exiting...\n");
            MemChkHelper::getInstance().record();
            break;
        case SIGQUIT:
            printf("Received SIGQUIT, cleaning up and exiting...\n");
            MemChkHelper::getInstance().record();
            break;
        case SIGABRT:
            printf("Received SIGABRT, cleaning up and exiting...\n");
            MemChkHelper::getInstance().record();
            break;
        case SIGILL:
            printf("Received SIGILL, cleaning up and exiting...\n");
            MemChkHelper::getInstance().record();
            break;
        case SIGFPE:
            printf("Received SIGFPE, cleaning up and exiting...\n");
            MemChkHelper::getInstance().record();
            break;
        case SIGSEGV:
            printf("Received SIGSEGV, cleaning up and exiting...\n");
            MemChkHelper::getInstance().record();
            break;
        default:
            break;
    }
    // 执行清理操作
    // 可以在这里添加释放资源、关闭文件等操作

    // 退出进程
    exit(EXIT_SUCCESS);
}

char** getbacktrace()
{
    char** bt = static_cast<char**>(__libc_malloc((MAX_BT_DEPTH + 1) * sizeof(char*)));

    unw_cursor_t cursor;
    unw_context_t context;

    // Initialize cursor to current frame for local unwinding.
    unw_getcontext(&context);
    unw_init_local(&cursor, &context);

    // Unwind frames one by one, going up the frame stack.
    int index = 0;
    while (unw_step(&cursor) > 0 && index < MAX_BT_DEPTH)
    {
        char* sym = static_cast<char*>(__libc_malloc(MAX_SYMBOL_LEN));

        // 通过函数指针获取函数对应的名称
        if (unw_get_proc_name(&cursor, sym, MAX_SYMBOL_LEN, nullptr) != 0)
        {
            strcpy(sym, "??");
        }
        bt[index] = sym;
        ++index;
    }
    bt[index] = nullptr;

    return bt;
}

// malloc
static std::atomic_bool malloc_dl_flag(true);
static std::atomic_bool malloc_print_flag(true);
extern "C" void *malloc(size_t len)
{
    void *addr = __libc_malloc(len);

    if (malloc_print_flag.load() && malloc_dl_flag.load())
    {
        malloc_print_flag.store(false);
        printf("add malloc info.\n");
        if(g_isProcInit)// 防止把启动进程时的malloc记录进去
        {
            char** attr = static_cast<char**>(__libc_malloc((MAX_BT_DEPTH + 1) * sizeof(char*)));
            char* strLen = static_cast<char*>(__libc_malloc(10));
            snprintf(strLen, 10, "%ld", len);
            attr[0] = strLen;
            attr[1] = nullptr;
            MemChkHelper::getInstance().addMallocInfo(addr, getbacktrace(), attr);
        }
        malloc_print_flag.store(true);
    }

    return addr;
}

// free
static std::atomic_bool free_print_flag(true);
extern "C" void free(void *addr)
{
    if (free_print_flag.load() && malloc_dl_flag.load())
    {
        free_print_flag.store(false);
        printf("remove malloc info.\n");
        if(g_isProcInit)
        {
            MemChkHelper::getInstance().removeMallocInfo(addr);
        }
        free_print_flag.store(true);
    }

    __libc_free(addr);
}


// new, new[]
typedef void *(*hooked_new)(size_t);
static hooked_new gcc_new = nullptr;
static hooked_new gcc_new_a = nullptr;

void *operator new(size_t len)
{
    if (!gcc_new)
    {
        malloc_dl_flag.store(false);
        dlopen("libstdc++.so", RTLD_LAZY);
        gcc_new = (hooked_new)(dlsym(RTLD_NEXT, "_Znwm"));
        malloc_dl_flag.store(true);
    }
    
    void *addr = gcc_new(len);

    printf("add new info.\n");
    if(g_isProcInit)
    {
        char** attr = static_cast<char**>(__libc_malloc((MAX_BT_DEPTH + 1) * sizeof(char*)));
        char* strLen = static_cast<char*>(__libc_malloc(10));
        snprintf(strLen, 10, "%ld", len);
        attr[0] = strLen;
        attr[1] = nullptr;
        MemChkHelper::getInstance().addNewInfo(addr, getbacktrace(), attr);
    }

    return addr;
}

void *operator new[](size_t len)
{
    if (!gcc_new_a)
    {
        malloc_dl_flag.store(false);
        dlopen("libstdc++.so", RTLD_LAZY);
        gcc_new_a = (hooked_new)(dlsym(RTLD_NEXT, "_Znam"));
        malloc_dl_flag.store(true);
    }
    
    void *addr = gcc_new_a(len);

    printf("add new[] info.\n");
    if(g_isProcInit)
    {
        char** attr = static_cast<char**>(__libc_malloc((MAX_BT_DEPTH + 1) * sizeof(char*)));
        char* strLen = static_cast<char*>(__libc_malloc(10));
        snprintf(strLen, 10, "%ld", len);
        attr[0] = strLen;
        attr[1] = nullptr;
        MemChkHelper::getInstance().addNewInfo(addr, getbacktrace(), attr);
    }

    return addr;
}

// delete, delete[]
typedef void (*hooked_delete)(void *);
static hooked_delete gcc_delete = nullptr;
static hooked_delete gcc_delete_a = nullptr;

void operator delete(void *addr)
{
    printf("remove new info.\n");
    if(g_isProcInit)
    {
        MemChkHelper::getInstance().removeNewInfo(addr);
    }

    if (!gcc_delete)
    {
        malloc_dl_flag.store(false);
        dlopen("libstdc++.so", RTLD_LAZY);
        gcc_delete = (hooked_delete)(dlsym(RTLD_NEXT, "_ZdlPv"));
        malloc_dl_flag.store(true);
    }
    
    gcc_delete(addr);
}

void operator delete[](void *addr)
{
    printf("remove new[] info.\n");
    if(g_isProcInit)
    {
        MemChkHelper::getInstance().removeNewInfo(addr);
    }

    if (!gcc_delete_a)
    {
        malloc_dl_flag.store(false);
        dlopen("libstdc++.so", RTLD_LAZY);
        gcc_delete_a = (hooked_delete)(dlsym(RTLD_NEXT, "_ZdaPv"));
        malloc_dl_flag.store(true);
    }
    
    gcc_delete_a(addr);
}